# 动态NSGA-II算法
# src/algorithms/nsga2.py
import numpy as np
import random
from typing import List, Dict, Any, Tuple
from dataclasses import dataclass
from copy import deepcopy
import warnings

@dataclass
class Individual:
    """个体类"""
    weights: np.ndarray
    objectives: np.ndarray
    constraints: np.ndarray
    rank: int = 0
    crowding_distance: float = 0.0
    
    def __post_init__(self):
        # 确保权重归一化
        if np.sum(self.weights) > 0:
            self.weights = self.weights / np.sum(self.weights)

class DynamicNSGA2:
    """动态NSGA-II算法"""
    
    def __init__(self, population_size: int = 100, 
                 generations: int = 50,
                 crossover_prob: float = 0.9,
                 mutation_prob: float = 0.1,
                 eta_c: float = 20.0,
                 eta_m: float = 20.0):
        
        self.population_size = population_size
        self.generations = generations
        self.crossover_prob = crossover_prob
        self.mutation_prob = mutation_prob
        self.eta_c = eta_c  # 交叉分布指数
        self.eta_m = eta_m  # 变异分布指数
        
        # 记忆机制
        self.memory_archive = []
        self.max_memory_size = 50
        
        # 种群多样性
        self.migration_rate = 0.15
        
    def initialize_population(self, n_assets: int) -> List[Individual]:
        """初始化种群"""
        population = []
        
        for _ in range(self.population_size):
            # 随机生成权重
            weights = np.random.random(n_assets)
            weights = weights / np.sum(weights)  # 归一化
            
            individual = Individual(
                weights=weights,
                objectives=np.zeros(3),  # 收益、风险、ESG
                constraints=np.zeros(2)  # 约束违反程度
            )
            population.append(individual)
        
        return population
    
    def evaluate_individual(self, individual: Individual, 
                          returns: np.ndarray,
                          cov_matrix: np.ndarray,
                          esg_scores: np.ndarray,
                          constraints: Dict[str, float]) -> None:
        """评估个体"""
        weights = individual.weights
        
        # 计算目标函数
        # 1. 期望收益
        expected_return = np.dot(weights, returns)
        
        # 2. 风险 (负的Sharpe比率)
        portfolio_variance = np.dot(weights.T, np.dot(cov_matrix, weights))
        portfolio_std = np.sqrt(portfolio_variance)
        sharpe_ratio = expected_return / portfolio_std if portfolio_std > 0 else 0
        
        # 3. ESG评分
        esg_score = np.dot(weights, esg_scores)
        
        # 目标函数 (最大化收益和ESG，最小化风险)
        individual.objectives = np.array([
            expected_return,      # 最大化收益
            -portfolio_std,       # 最小化风险
            esg_score            # 最大化ESG
        ])
        
        # 约束违反
        # 1. 单一资产权重约束
        max_weight_violation = max(0, np.max(weights) - constraints['max_single_weight'])
        
        # 2. 现金比例约束
        cash_violation = max(0, constraints['min_cash_ratio'] - np.sum(weights))
        
        individual.constraints = np.array([max_weight_violation, cash_violation])
    
    def fast_non_dominated_sort(self, population: List[Individual]) -> List[List[int]]:
        """快速非支配排序"""
        n = len(population)
        domination_count = [0] * n
        dominated_solutions = [[] for _ in range(n)]
        fronts = [[]]
        
        # 计算支配关系
        for i in range(n):
            for j in range(n):
                if i != j:
                    if self.dominates(population[i], population[j]):
                        dominated_solutions[i].append(j)
                    elif self.dominates(population[j], population[i]):
                        domination_count[i] += 1
            
            if domination_count[i] == 0:
                population[i].rank = 0
                fronts[0].append(i)
        
        # 构建后续前沿
        front_idx = 0
        try:
            while len(fronts[front_idx]) > 0 and front_idx < len(fronts) - 1:
                next_front = []
                for i in fronts[front_idx]:
                    for j in dominated_solutions[i]:
                        domination_count[j] -= 1
                        if domination_count[j] == 0:
                            population[j].rank = front_idx + 1
                            next_front.append(j)
                
                front_idx += 1
                if len(next_front) > 0:
                    fronts.append(next_front)
        except Exception as e:
            print(e)
        
        return fronts
    
    def dominates(self, ind1: Individual, ind2: Individual) -> bool:
        """判断ind1是否支配ind2"""
        # 考虑约束
        if np.sum(ind1.constraints) > 0 and np.sum(ind2.constraints) > 0:
            return np.sum(ind1.constraints) < np.sum(ind2.constraints)
        elif np.sum(ind1.constraints) > 0:
            return False
        elif np.sum(ind2.constraints) > 0:
            return True
        
        # 无约束违反时比较目标函数
        better_in_any = False
        for i in range(len(ind1.objectives)):
            if ind1.objectives[i] < ind2.objectives[i]:
                return False
            elif ind1.objectives[i] > ind2.objectives[i]:
                better_in_any = True
        
        return better_in_any
    
    def calculate_crowding_distance(self, front: List[int], 
                                  population: List[Individual]) -> None:
        """计算拥挤距离"""
        if len(front) <= 2:
            for i in front:
                population[i].crowding_distance = float('inf')
            return
        
        # 初始化拥挤距离
        for i in front:
            population[i].crowding_distance = 0
        
        n_objectives = len(population[0].objectives)
        
        # 对每个目标函数计算拥挤距离
        for obj_idx in range(n_objectives):
            # 按目标函数值排序
            front.sort(key=lambda x: population[x].objectives[obj_idx])
            
            # 边界个体设为无穷大
            population[front[0]].crowding_distance = float('inf')
            population[front[-1]].crowding_distance = float('inf')
            
            # 计算中间个体的拥挤距离
            obj_min = population[front[0]].objectives[obj_idx]
            obj_max = population[front[-1]].objectives[obj_idx]
            
            if obj_max - obj_min > 0:
                for i in range(1, len(front) - 1):
                    distance = (population[front[i+1]].objectives[obj_idx] - 
                              population[front[i-1]].objectives[obj_idx]) / (obj_max - obj_min)
                    population[front[i]].crowding_distance += distance
    
    def tournament_selection(self, population: List[Individual], 
                           tournament_size: int = 2) -> Individual:
        """锦标赛选择"""
        tournament = random.sample(population, tournament_size)
        
        # 选择最好的个体
        best = tournament[0]
        for individual in tournament[1:]:
            if (individual.rank < best.rank or 
                (individual.rank == best.rank and 
                 individual.crowding_distance > best.crowding_distance)):
                best = individual
        
        return deepcopy(best)
    
    def sbx_crossover(self, parent1: Individual, parent2: Individual) -> Tuple[Individual, Individual]:
        """SBX交叉"""
        n_vars = len(parent1.weights)
        child1_weights = np.zeros(n_vars)
        child2_weights = np.zeros(n_vars)
        
        for i in range(n_vars):
            if random.random() <= 0.5:
                y1, y2 = parent1.weights[i], parent2.weights[i]
                
                if abs(y1 - y2) > 1e-14:
                    if y1 > y2:
                        y1, y2 = y2, y1
                    
                    # 计算βq
                    rand = random.random()
                    if rand <= 0.5:
                        beta = (2 * rand) ** (1 / (self.eta_c + 1))
                    else:
                        beta = (1 / (2 * (1 - rand))) ** (1 / (self.eta_c + 1))
                    
                    # 生成子代
                    child1_weights[i] = 0.5 * ((y1 + y2) - beta * (y2 - y1))
                    child2_weights[i] = 0.5 * ((y1 + y2) + beta * (y2 - y1))
                else:
                    child1_weights[i] = y1
                    child2_weights[i] = y2
            else:
                child1_weights[i] = parent1.weights[i]
                child2_weights[i] = parent2.weights[i]
        
        # 确保权重为正且归一化
        child1_weights = np.maximum(child1_weights, 0)
        child2_weights = np.maximum(child2_weights, 0)
        
        if np.sum(child1_weights) > 0:
            child1_weights = child1_weights / np.sum(child1_weights)
        if np.sum(child2_weights) > 0:
            child2_weights = child2_weights / np.sum(child2_weights)
        
        child1 = Individual(weights=child1_weights, objectives=np.zeros(3), constraints=np.zeros(2))
        child2 = Individual(weights=child2_weights, objectives=np.zeros(3), constraints=np.zeros(2))
        
        return child1, child2
    
    def polynomial_mutation(self, individual: Individual) -> Individual:
        """多项式变异"""
        mutated_weights = individual.weights.copy()
        
        for i in range(len(mutated_weights)):
            if random.random() <= self.mutation_prob:
                y = mutated_weights[i]
                delta1 = y
                delta2 = 1 - y
                
                mut_pow = 1.0 / (self.eta_m + 1.0)
                rand = random.random()
                
                if rand <= 0.5:
                    xy = 1.0 - delta1
                    val = 2.0 * rand + (1.0 - 2.0 * rand) * (xy ** (self.eta_m + 1.0))
                    deltaq = (val ** mut_pow) - 1.0
                else:
                    xy = 1.0 - delta2
                    val = 2.0 * (1.0 - rand) + 2.0 * (rand - 0.5) * (xy ** (self.eta_m + 1.0))
                    deltaq = 1.0 - (val ** mut_pow)
                
                mutated_weights[i] = y + deltaq
        
        # 确保权重为正且归一化
        mutated_weights = np.maximum(mutated_weights, 0)
        if np.sum(mutated_weights) > 0:
            mutated_weights = mutated_weights / np.sum(mutated_weights)
        
        return Individual(weights=mutated_weights, objectives=np.zeros(3), constraints=np.zeros(2))
    
    def environmental_selection(self, population: List[Individual]) -> List[Individual]:
        """环境选择"""
        # 非支配排序
        fronts = self.fast_non_dominated_sort(population)
        
        # 计算拥挤距离
        for front in fronts:
            self.calculate_crowding_distance(front, population)
        
        # 选择个体
        selected = []
        for front in fronts:
            if len(selected) + len(front) <= self.population_size:
                selected.extend([population[i] for i in front])
            else:
                # 按拥挤距离排序
                remaining = self.population_size - len(selected)
                front_individuals = [(population[i], i) for i in front]
                front_individuals.sort(key=lambda x: x[0].crowding_distance, reverse=True)
                
                for i in range(remaining):
                    selected.append(front_individuals[i][0])
                break
        
        return selected
    
    def update_memory(self, population: List[Individual]) -> None:
        """更新记忆存档"""
        # 获取第一前沿
        fronts = self.fast_non_dominated_sort(population)
        if len(fronts) > 0:
            first_front = [population[i] for i in fronts[0]]
            
            # 添加到记忆存档
            self.memory_archive.extend(first_front)
            
            # 限制存档大小
            if len(self.memory_archive) > self.max_memory_size:
                self.memory_archive = self.memory_archive[-self.max_memory_size:]
    
    def inject_memory(self, population: List[Individual]) -> List[Individual]:
        """注入记忆中的解"""
        if len(self.memory_archive) == 0:
            return population
        
        # 选择一部分个体替换
        n_replace = int(self.population_size * self.migration_rate)
        
        # 随机选择要替换的个体
        replace_indices = random.sample(range(len(population)), n_replace)
        
        # 从记忆存档中随机选择解
        memory_solutions = random.sample(self.memory_archive, 
                                       min(n_replace, len(self.memory_archive)))
        
        # 替换
        for i, idx in enumerate(replace_indices):
            if i < len(memory_solutions):
                population[idx] = deepcopy(memory_solutions[i])
        
        return population
    
    def optimize(self, returns: np.ndarray, cov_matrix: np.ndarray,
                esg_scores: np.ndarray, constraints: Dict[str, float],
                market_changed: bool = False) -> List[Individual]:
        """主优化函数"""
        n_assets = len(returns)
        
        # 初始化种群
        population = self.initialize_population(n_assets)
        
        # 如果检测到市场变化，注入记忆中的解
        if market_changed:
            population = self.inject_memory(population)
        
        # 评估初始种群
        for individual in population:
            self.evaluate_individual(individual, returns, cov_matrix, esg_scores, constraints)
        
        # 进化迭代
        for generation in range(self.generations):
            # 生成子代
            offspring = []
            
            for _ in range(self.population_size // 2):
                # 选择父代
                parent1 = self.tournament_selection(population)
                parent2 = self.tournament_selection(population)
                
                # 交叉
                if random.random() <= self.crossover_prob:
                    child1, child2 = self.sbx_crossover(parent1, parent2)
                else:
                    child1, child2 = deepcopy(parent1), deepcopy(parent2)
                
                # 变异
                child1 = self.polynomial_mutation(child1)
                child2 = self.polynomial_mutation(child2)
                
                # 评估子代
                self.evaluate_individual(child1, returns, cov_matrix, esg_scores, constraints)
                self.evaluate_individual(child2, returns, cov_matrix, esg_scores, constraints)
                
                offspring.extend([child1, child2])
            
            # 合并父代和子代
            combined_population = population + offspring
            
            # 环境选择
            population = self.environmental_selection(combined_population)
            
            # 更新记忆存档
            if generation % 10 == 0:
                self.update_memory(population)
        
        return population
